"""
页式存储系统 V2.0
完全符合实践要求：4KB页面、分配释放、读写操作
"""

import os
import struct
import json
from typing import Dict, List, Optional, Any, Tuple
import time

class Page:
    """页面类"""
    
    def __init__(self, page_id: int, data: bytearray = None):
        """初始化页面"""
        self.page_id = page_id
        self.size = 4096  # 4KB
        self.data = data if data else bytearray(self.size)
        self.dirty = False
        self.last_access = time.time()
    
    def read_data(self, offset: int, length: int) -> bytearray:
        """读取页面数据"""
        self.last_access = time.time()
        return self.data[offset:offset + length]
    
    def write_data(self, offset: int, data: bytearray):
        """写入页面数据"""
        self.data[offset:offset + len(data)] = data
        self.dirty = True
        self.last_access = time.time()
    
    def is_full(self) -> bool:
        """检查页面是否已满"""
        return len([b for b in self.data if b != 0]) >= self.size * 0.9  # 90%满
    
    def get_free_space(self) -> int:
        """获取空闲空间"""
        used = len([b for b in self.data if b != 0])
        return self.size - used

class PageStorageSystem:
    """页式存储系统"""
    
    def __init__(self, data_dir: str = "data"):
        """初始化页式存储系统"""
        self.data_dir = data_dir
        self.page_size = 4096  # 4KB
        self.next_page_id = 0
        self.free_pages = []  # 空闲页面列表
        self.allocated_pages = set()  # 已分配页面集合
        
        # 确保数据目录存在
        os.makedirs(data_dir, exist_ok=True)
        
        # 加载页面分配信息
        self.load_page_info()
    
    def load_page_info(self):
        """加载页面分配信息"""
        page_info_file = os.path.join(self.data_dir, "page_info.json")
        if os.path.exists(page_info_file):
            try:
                with open(page_info_file, 'r', encoding='utf-8') as f:
                    info = json.load(f)
                    self.next_page_id = info.get('next_page_id', 0)
                    self.free_pages = info.get('free_pages', [])
                    self.allocated_pages = set(info.get('allocated_pages', []))
            except Exception as e:
                print(f"加载页面信息失败: {e}")
    
    def save_page_info(self):
        """保存页面分配信息"""
        page_info_file = os.path.join(self.data_dir, "page_info.json")
        try:
            info = {
                'next_page_id': self.next_page_id,
                'free_pages': self.free_pages,
                'allocated_pages': list(self.allocated_pages)
            }
            with open(page_info_file, 'w', encoding='utf-8') as f:
                json.dump(info, f, indent=2)
        except Exception as e:
            print(f"保存页面信息失败: {e}")
    
    def allocate_page(self) -> int:
        """分配新页面"""
        if self.free_pages:
            # 使用空闲页面
            page_id = self.free_pages.pop()
        else:
            # 创建新页面
            page_id = self.next_page_id
            self.next_page_id += 1
        
        self.allocated_pages.add(page_id)
        self.save_page_info()
        
        # 创建页面文件
        self.write_page(page_id, bytearray(self.page_size))
        
        return page_id
    
    def deallocate_page(self, page_id: int):
        """释放页面"""
        if page_id in self.allocated_pages:
            self.allocated_pages.remove(page_id)
            self.free_pages.append(page_id)
            self.save_page_info()
            
            # 删除页面文件
            page_file = self.get_page_file_path(page_id)
            if os.path.exists(page_file):
                os.remove(page_file)
    
    def read_page(self, page_id: int) -> Optional[bytearray]:
        """读取页面数据"""
        page_file = self.get_page_file_path(page_id)
        
        try:
            with open(page_file, 'rb') as f:
                data = f.read(self.page_size)
                if len(data) == self.page_size:
                    return bytearray(data)
                else:
                    # 页面大小不正确，返回空页面
                    return bytearray(self.page_size)
        except FileNotFoundError:
            # 页面文件不存在，返回空页面
            return bytearray(self.page_size)
        except Exception as e:
            print(f"读取页面 {page_id} 失败: {e}")
            return None
    
    def write_page(self, page_id: int, data: bytearray) -> bool:
        """写入页面数据"""
        if len(data) != self.page_size:
            print(f"页面大小不正确: {len(data)} != {self.page_size}")
            return False
        
        page_file = self.get_page_file_path(page_id)
        
        try:
            with open(page_file, 'wb') as f:
                f.write(data)
            return True
        except Exception as e:
            print(f"写入页面 {page_id} 失败: {e}")
            return False
    
    def get_page_file_path(self, page_id: int) -> str:
        """获取页面文件路径"""
        return os.path.join(self.data_dir, f"page_{page_id:06d}.dat")
    
    def get_page_count(self) -> int:
        """获取总页面数"""
        return len(self.allocated_pages)
    
    def get_free_page_count(self) -> int:
        """获取空闲页面数"""
        return len(self.free_pages)
    
    def get_allocated_page_count(self) -> int:
        """获取已分配页面数"""
        return len(self.allocated_pages)
    
    def print_page_info(self):
        """打印页面信息"""
        print("页式存储系统信息:")
        print(f"  页面大小: {self.page_size} 字节")
        print(f"  总页面数: {self.get_page_count()}")
        print(f"  已分配页面数: {self.get_allocated_page_count()}")
        print(f"  空闲页面数: {self.get_free_page_count()}")
        print(f"  下一个页面ID: {self.next_page_id}")
        
        if self.allocated_pages:
            print("  已分配页面:")
            for page_id in sorted(self.allocated_pages):
                print(f"    {page_id}")

class TablePageManager:
    """表页面管理器"""
    
    def __init__(self, storage_system: PageStorageSystem):
        """初始化表页面管理器"""
        self.storage = storage_system
        self.table_pages = {}  # 表名 -> 页面ID列表
    
    def create_table(self, table_name: str) -> bool:
        """为表创建初始页面"""
        if table_name in self.table_pages:
            return False
        
        # 分配初始页面
        page_id = self.storage.allocate_page()
        self.table_pages[table_name] = [page_id]
        
        # 初始化页面头部
        page_data = self.storage.read_page(page_id)
        if page_data:
            # 页面头部: [记录数(4字节)][空闲空间(4字节)][数据...]
            struct.pack_into('<I', page_data, 0, 0)  # 记录数
            struct.pack_into('<I', page_data, 4, self.storage.page_size - 8)  # 空闲空间
            self.storage.write_page(page_id, page_data)
        
        return True
    
    def drop_table(self, table_name: str) -> bool:
        """删除表的所有页面"""
        if table_name not in self.table_pages:
            return False
        
        # 释放所有页面
        for page_id in self.table_pages[table_name]:
            self.storage.deallocate_page(page_id)
        
        del self.table_pages[table_name]
        return True
    
    def get_table_pages(self, table_name: str) -> List[int]:
        """获取表的所有页面"""
        return self.table_pages.get(table_name, [])
    
    def add_page_to_table(self, table_name: str) -> int:
        """为表添加新页面"""
        if table_name not in self.table_pages:
            return -1
        
        page_id = self.storage.allocate_page()
        self.table_pages[table_name].append(page_id)
        
        # 初始化页面头部
        page_data = self.storage.read_page(page_id)
        if page_data:
            struct.pack_into('<I', page_data, 0, 0)  # 记录数
            struct.pack_into('<I', page_data, 4, self.storage.page_size - 8)  # 空闲空间
            self.storage.write_page(page_id, page_data)
        
        return page_id
    
    def get_page_record_count(self, page_id: int) -> int:
        """获取页面的记录数"""
        page_data = self.storage.read_page(page_id)
        if page_data:
            return struct.unpack('<I', page_data[0:4])[0]
        return 0
    
    def set_page_record_count(self, page_id: int, count: int):
        """设置页面的记录数"""
        page_data = self.storage.read_page(page_id)
        if page_data:
            struct.pack_into('<I', page_data, 0, count)
            self.storage.write_page(page_id, page_data)
    
    def get_page_free_space(self, page_id: int) -> int:
        """获取页面的空闲空间"""
        page_data = self.storage.read_page(page_id)
        if page_data:
            return struct.unpack('<I', page_data[4:8])[0]
        return 0
    
    def set_page_free_space(self, page_id: int, free_space: int):
        """设置页面的空闲空间"""
        page_data = self.storage.read_page(page_id)
        if page_data:
            struct.pack_into('<I', page_data, 4, free_space)
            self.storage.write_page(page_id, page_data)

# 测试函数
if __name__ == "__main__":
    storage = PageStorageSystem("test_storage")
    table_manager = TablePageManager(storage)
    
    print("测试页式存储系统...")
    
    # 测试页面分配
    print("\n1. 测试页面分配...")
    page_id1 = storage.allocate_page()
    page_id2 = storage.allocate_page()
    print(f"分配页面: {page_id1}, {page_id2}")
    
    # 测试页面读写
    print("\n2. 测试页面读写...")
    test_data = bytearray(4096)
    test_data[0:4] = b'TEST'
    storage.write_page(page_id1, test_data)
    
    read_data = storage.read_page(page_id1)
    print(f"写入数据: {test_data[0:4]}")
    print(f"读取数据: {read_data[0:4]}")
    
    # 测试表页面管理
    print("\n3. 测试表页面管理...")
    table_manager.create_table("test_table")
    pages = table_manager.get_table_pages("test_table")
    print(f"表页面: {pages}")
    
    # 打印存储信息
    print("\n4. 存储系统信息:")
    storage.print_page_info()
    
    # 清理
    storage.deallocate_page(page_id1)
    storage.deallocate_page(page_id2)
    table_manager.drop_table("test_table")
